import os
from zhipuai import ZhipuAI


def get_api_key():
    """
    获取API密钥。

    Returns:
        str: API密钥。

    Raises:
        ValueError: 如果API密钥未找到。
    """
    key = os.getenv("ZHIPUAI_API_KEY")
    if not key:
        raise ValueError("API密钥未找到。请确保已设置环境变量ZHIPUAI_API_KEY。")
    return key


def initialize_client(api_key):
    """
    初始化ZhipuAI客户端。

    Args:
        api_key (str): API密钥。

    Returns:
        ZhipuAI: 初始化的ZhipuAI客户端。
    """
    return ZhipuAI(api_key=api_key)


def create_completion(client, messages):
    """
    创建聊天完成请求。

    Args:
        client (ZhipuAI): ZhipuAI客户端。
        messages (list): 对话消息列表。

    Returns:
        response: API响应。

    Raises:
        Exception: 如果API调用失败。
    """
    try:
        response = client.chat.completions.create(
            model="glm-4-flash",   # 模型名称
            messages=messages,     # 对话消息列表
            stream=True            # 流式输出
        )
        return response
    except Exception as e:
        raise Exception(f"API调用失败: {e}")


def stream_response(response):
    """
    流式输出响应内容。

    Args:
        response: API响应。

    Returns:
        str: 助手的完整回复。
    """
    ai_response = ""
    for chunk in response:
        if chunk.choices and chunk.choices[0].delta and chunk.choices[0].delta.content:
            print(chunk.choices[0].delta.content, end='', flush=True)
            ai_response += chunk.choices[0].delta.content
    print()
    return ai_response


def main():
    """
    主函数。
    """
    api_key = get_api_key()
    client = initialize_client(api_key)
    messages = [
        {"role": "system", "content": "你是一个有趣的人类，你的名字叫做小明"},
    ]

    while True:
        user_input = input("您: ")
        if user_input.lower() in ["exit", "quit"]:
            break
        messages.append({"role": "user", "content": user_input})
        response = create_completion(client, messages)
        assistant_response = stream_response(response)
        messages.append({"role": "assistant", "content": assistant_response})  # 添加助手的回复到历史记录


if __name__ == "__main__":
    main()
